import os
os.environ["OMP_NUM_THREADS"] = "1"        # 控制 OpenMP 线程数
os.environ["MKL_NUM_THREADS"] = "1"       # 控制 Intel MKL 线程数
os.environ["OPENBLAS_NUM_THREADS"] = "1"  # 控制 OpenBLAS 线程数
os.environ["NUMEXPR_NUM_THREADS"] = "1"   # 控制 NumExpr 线程数
import numpy as np
from cffi import FFI
from broadcast_add import broadcast_add_wrapper


# 创建测试数据（支持任意维度）
a = np.random.uniform(-10, 10, (10, 500, 3000))
b = np.random.uniform(-10, 10, (10, 500, 3000))

# 验证结果
expected = broadcast_add_wrapper(a, b, 'np')  # 使用Numpy原生广播
# 执行广播加法
result = broadcast_add_wrapper(a, b, 'c')

print(np.sum(np.abs(result - expected)))
